# -*- coding: utf-8 -*-
"""
all_single_guard_modbus_fc10_horn_then_hex.py

功能：
1) 先生成 Horn（Datalog/CHC）单守卫“违反集合”规则文件：./horn_rules_modbus_fc10.dl
2) 再生成逐守卫的 HEX 数据包文件（每行附带属性说明），存放到 ./out_labeled/<guard>.hexlabeled

说明：
- 仅覆盖 Modbus/TCP 的 FC=0x10（除 G3 故意换 FC），遵循你给的最小从站模型。
- 为避免 G5 全集枚举导致卡顿，提供了“采样配置” CFG；你可把范围改大至“全集”。
- 每条 HEX 行末尾均附上：pid,len,fc,start,qty,bc,vbytes 的属性值，便于溯源。
"""

from pathlib import Path
from typing import Iterable, Tuple, Dict, Any, List
from dataclasses import dataclass

# ===================== 采样配置（可放大到“全集”） =====================
CFG = {
    # G1: pid != 0 的全集可改成 range(1, 65536)
    "G1_pid_values": list(range(1, 11)),
    # G3: 非 0x10 的功能码集合（可扩充到除 0x10 外的 0..255）
    "G3_fc_values":  [0x00, 0x01, 0x02, 0x03, 0x11, 0x2B, 0x90],
    # G4: PDU 长度 {1..5}
    "G4_vbytes":     [1, 2, 3, 4, 5],
    # G5: 指定若干 qty，并在等式两侧取多个 delta（避免 delta=0）
    "G5_qty_values": [0, 1, 3, 10, 100],
    "G5_deltas":     [-2, -1, 1, 2],
    # G6: 低侧（qty=0）与高侧（qty>=124）示例
    "G6_qty_low":    [0],
    "G6_qty_high":   [124, 125, 200, 1000],
    # G7: 在 1..123 内挑几个 qty，并取一组 bc 值（自动过滤等于 2*qty 的点）
    "G7_qty_values": [1, 2, 3, 10, 50, 100, 123],
    "G7_bc_pool":    list(range(0, 8)) + [0xFE, 0xFF],
    # G8: 几个 qty；start 从越界起点开始取 span 个
    "G8_qty_values": [1, 2, 3, 10, 50, 123],
    "G8_start_span": 5,
}

# ===================== 常量与工具 =====================
REG_MAX = 2000
HOST_UID = 0x01
TID = 0x0001
PID_OK = 0x0000

def u16(x: int) -> bytes:
    return x.to_bytes(2, "big", signed=False)

def u8(x: int) -> bytes:
    return x.to_bytes(1, "big", signed=False)

def hexline(b: bytes) -> str:
    return " ".join(f"{v:02X}" for v in b)

@dataclass
class Pkt:
    pid: int
    length: int   # MBAP LEN（= 1 + vbytes）
    fc: int
    start: int
    qty: int
    bc: int
    vbytes: int   # PDU size

def build_mbap_pdu(pkt: Pkt, values_len: int, fc_override: int = None, raw_pdu: bytes = None) -> bytes:
    """
    组装 MBAP+PDU；
    - raw_pdu：用于 G4 直接给“短头部”的 PDU；
    - fc_override：用于 G3 改写 fc。
    """
    # 构造 PDU
    if raw_pdu is not None:
        pdu = raw_pdu
    else:
        pdu_fc = pkt.fc if fc_override is None else fc_override
        if pdu_fc != 0x10:
            pdu = u8(pdu_fc) + bytes(max(0, pkt.vbytes - 1))
        else:
            header = u8(0x10) + u16(pkt.start) + u16(pkt.qty) + u8(pkt.bc)  # 6 bytes
            pad = bytes([0x00, 0x01]) * (max(0, values_len) // 2) + (b"\\x00" if values_len % 2 else b"")
            pdu = header + pad[:values_len]
    # 一致性检查
    assert pkt.length == 1 + len(pdu), f"LEN mismatch: {pkt.length} vs {1 + len(pdu)}"
    # MBAP
    mbap = u16(TID) + u16(pkt.pid) + u16(pkt.length) + u8(HOST_UID)
    return mbap + pdu

def label(attrs: Dict[str, Any]) -> str:
    order = ["pid","len","fc","start","qty","bc","vbytes"]
    return " ".join([f"{k}={attrs[k]}" for k in order if k in attrs])

def write_labeled(path: Path, rows: Iterable[Tuple[bytes, Dict[str, Any]]]) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    n = 0
    with path.open("w", encoding="utf-8") as f:
        for raw, attrs in rows:
            f.write(f"{hexline(raw)}  # {label(attrs)}\\n")
            n += 1
    return n

# ===================== Horn 规则（“违反集合”） =====================
HORN = r"""// horn_rules_modbus_fc10.dl (single-guard violations, no enumeration)
.decl DomPID(pid:int) .decl DomLEN(len:int) .decl DomFC(fc:int) .decl DomStart(start:int)
.decl DomQty(qty:int) .decl DomBC(bc:int)  .decl DomVBytes(v:int)
.decl LenOK(len:int, qty:int)   .decl QtyOK(qty:int)   .decl FcOK(fc:int)
.decl BcOK(bc:int, qty:int)     .decl AddrOK(start:int, qty:int)
LenOK(len,qty) :- len = 7 + 2*qty.
QtyOK(qty)    :- qty >= 1, qty <= 123.
FcOK(fc)      :- fc = 16.
BcOK(bc,qty)  :- bc = 2*qty.
AddrOK(s,q)   :- s + q <= 2000.

.decl ViolG1(pid:int, len:int)
.decl ViolG2(len:int)
.decl ViolG3(fc:int, len:int)
.decl ViolG4(vbytes:int, len:int)
.decl ViolG5_low(qty:int, len:int)
.decl ViolG5_high(qty:int, len:int)
.decl ViolG6_low(qty:int, len:int)
.decl ViolG6_high(qty:int, len:int)
.decl ViolG7(qty:int, bc:int, len:int)
.decl ViolG8(start:int, qty:int, len:int)

ViolG1(pid,len) :- DomPID(pid), DomLEN(len), pid != 0, len >= 1.
ViolG2(len)     :- DomLEN(len), len = 0.
ViolG3(fc,len)  :- DomFC(fc), DomLEN(len), len >= 2, fc != 16.
ViolG4(v,len)   :- DomVBytes(v), DomLEN(len), len = 1 + v, v >= 1, FcOK(16), v < 6.
ViolG5_low(q,l) :- DomQty(q), DomLEN(l), DomVBytes(v), l = 1 + v, v >= 6, FcOK(16), l < 7 + 2*q.
ViolG5_high(q,l):- DomQty(q), DomLEN(l), DomVBytes(v), l = 1 + v, v >= 6, FcOK(16), l > 7 + 2*q.
ViolG6_low(q,l) :- DomQty(q), DomLEN(l), DomVBytes(v), q = 0, l = 7 + 2*q, v = l - 1, v >= 6, FcOK(16).
ViolG6_high(q,l):- DomQty(q), DomLEN(l), DomVBytes(v), q >= 124, l = 7 + 2*q, v = l - 1, v >= 6, FcOK(16).
ViolG7(q,bc,l)  :- DomQty(q), DomBC(bc), DomLEN(l), DomVBytes(v), QtyOK(q), l = 7 + 2*q, v = l - 1, v >= 6, FcOK(16), bc != 2*q.
ViolG8(s,q,l)   :- DomStart(s), DomQty(q), DomLEN(l), DomVBytes(v), QtyOK(q), l = 7 + 2*q, v = l - 1, v >= 6, FcOK(16), BcOK(2*q,q), s + q > 2000.
""".strip()

def write_horn_file(path: Path) -> None:
    path.write_text(HORN, encoding="utf-8")

# ===================== 逐守卫生成（带属性标签） =====================

def gen_G1_rows() -> Iterable[Tuple[bytes, Dict[str, Any]]]:
    # 固定规范 PDU（Qty=1, BC=2, Values=2B）→ vbytes=8, len=9
    qty = 1; bc = 2; start = 0
    vbytes = 1 + 2 + 2 + 1 + 2
    length = 1 + vbytes
    for pid in CFG["G1_pid_values"]:
        pkt = Pkt(pid=pid, length=length, fc=0x10, start=start, qty=qty, bc=bc, vbytes=vbytes)
        raw = build_mbap_pdu(pkt, values_len=2)
        yield raw, {"pid":pid,"len":length,"fc":0x10,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes}

def gen_G2_rows() -> Iterable[Tuple[bytes, Dict[str, Any]]]:
    # LEN=0；无 PDU
    pid=PID_OK; length=0; fc=0; start=0; qty=0; bc=0; vbytes=0
    mbap = u16(TID) + u16(pid) + u16(length) + u8(HOST_UID)
    yield mbap, {"pid":pid,"len":length,"fc":fc,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes}

def gen_G3_rows() -> Iterable[Tuple[bytes, Dict[str, Any]]]:
    # len=2, vbytes=1（仅 1 字节 fc）
    length=2; vbytes=1; pid=PID_OK; start=0; qty=0; bc=0
    for fc in CFG["G3_fc_values"]:
        pkt = Pkt(pid=pid, length=length, fc=fc, start=start, qty=qty, bc=bc, vbytes=vbytes)
        raw = build_mbap_pdu(pkt, values_len=0, fc_override=fc)
        yield raw, {"pid":pid,"len":length,"fc":fc,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes}

def gen_G4_rows() -> Iterable[Tuple[bytes, Dict[str, Any]]]:
    # 直接 raw_pdu：0x10 + 零填充，使 vbytes ∈ {1..5}
    pid=PID_OK; fc=0x10; start=0; qty=0; bc=0
    for vbytes in CFG["G4_vbytes"]:
        length = 1 + vbytes
        raw_pdu = bytes([0x10] + [0x00]*(vbytes-1)) if vbytes >= 1 else b""
        pkt = Pkt(pid=pid, length=length, fc=fc, start=start, qty=qty, bc=bc, vbytes=vbytes)
        raw = build_mbap_pdu(pkt, values_len=0, raw_pdu=raw_pdu)
        yield raw, {"pid":pid,"len":length,"fc":fc,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes}

def gen_G5_rows() -> Iterable[Tuple[bytes, Dict[str, Any]]]:
    # 违反 len = 7 + 2*qty，两侧各采几个 delta
    pid=PID_OK; fc=0x10; start=0; bc=0
    for qty in CFG["G5_qty_values"]:
        base = 7 + 2*qty
        for delta in CFG["G5_deltas"]:
            if delta == 0:
                continue
            length = base + delta
            if not (7 <= length <= 0xFFFF):
                continue
            vbytes = length - 1
            header = 6
            values_len = max(0, vbytes - header)
            pkt = Pkt(pid=pid, length=length, fc=fc, start=start, qty=qty, bc=bc, vbytes=vbytes)
            raw = build_mbap_pdu(pkt, values_len=values_len)
            which = "low" if length < base else "high"
            yield raw, {"pid":pid,"len":length,"fc":fc,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes,"G5_side":which}

def gen_G6_rows() -> Iterable[Tuple[bytes, Dict[str, Any]]]:
    # qty 下溢与上溢（注意：保持 len = 7 + 2*qty，让读取流程能进入该守卫）
    pid=PID_OK; fc=0x10; start=0
    # 低侧 qty=0
    qty=0; length=7+2*qty; vbytes=length-1; bc=0
    pkt = Pkt(pid=pid, length=length, fc=fc, start=start, qty=qty, bc=bc, vbytes=vbytes)
    raw = build_mbap_pdu(pkt, values_len=max(0, vbytes-6))
    yield raw, {"pid":pid,"len":length,"fc":fc,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes,"G6_side":"low"}
    # 高侧 qty≥124
    for qty in CFG["G6_qty_high"]:
        length = 7 + 2*qty
        if length > 0xFFFF:
            continue
        vbytes = length - 1
        bc = (2*qty) & 0xFF   # 让后续检查（若有）通过 G7
        pkt = Pkt(pid=pid, length=length, fc=fc, start=start, qty=qty, bc=bc, vbytes=vbytes)
        raw = build_mbap_pdu(pkt, values_len=max(0, vbytes-6))
        yield raw, {"pid":pid,"len":length,"fc":fc,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes,"G6_side":"high"}

def gen_G7_rows() -> Iterable[Tuple[bytes, Dict[str, Any]]]:
    # 在通过 G5/G6 的前提下（len=7+2*qty 且 1<=qty<=123），选择 bc != 2*qty
    pid=PID_OK; fc=0x10; start=0
    for qty in CFG["G7_qty_values"]:
        length = 7 + 2*qty
        vbytes = length - 1
        values_len = 2*qty
        for bc in CFG["G7_bc_pool"]:
            if bc == 2*qty:
                continue
            pkt = Pkt(pid=pid, length=length, fc=fc, start=start, qty=qty, bc=bc, vbytes=vbytes)
            raw = build_mbap_pdu(pkt, values_len=values_len)
            yield raw, {"pid":pid,"len":length,"fc":fc,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes}

def gen_G8_rows() -> Iterable[Tuple[bytes, Dict[str, Any]]]:
    # 在通过 G5/G6/G7 的前提下，令 start + qty > 2000
    pid=PID_OK; fc=0x10
    for qty in CFG["G8_qty_values"]:
        length = 7 + 2*qty
        if length > 0xFFFF:
            continue
        vbytes = length - 1
        bc = 2*qty
        start0 = max(0, 2001 - qty)  # 越界起点
        for i in range(CFG["G8_start_span"]):
            start = start0 + i
            values_len = 2*qty
            pkt = Pkt(pid=pid, length=length, fc=fc, start=start, qty=qty, bc=bc, vbytes=vbytes)
            raw = build_mbap_pdu(pkt, values_len=values_len)
            yield raw, {"pid":pid,"len":length,"fc":fc,"start":start,"qty":qty,"bc":bc,"vbytes":vbytes}

# ===================== 主流程 =====================
def main():
    base = Path(".")
    # 1) 写 Horn 规则
    horn_path = base / "horn_rules_modbus_fc10.dl"
    horn_path.write_text(HORN, encoding="utf-8")
    # 2) 生成带标签的 HEX
    outdir = base / "out_labeled"
    outdir.mkdir(parents=True, exist_ok=True)
    (outdir / "README.txt").write_text(
        "本目录为“单守卫反例（带属性标签）”的可控采样版本。\\n"
        "你可在脚本 CFG 中放大采样范围以接近/达到全集；当前配置：\\n\\n"
        + "\\n".join([f"{k}: {v}" for k,v in CFG.items()]) + "\\n",
        encoding="utf-8"
    )
    specs = [
        ("G1_PID_NEQ_0.hexlabeled", gen_G1_rows),
        ("G2_LEN_EQ_0.hexlabeled", gen_G2_rows),
        ("G3_FC_NEQ_10.hexlabeled", gen_G3_rows),
        ("G4_PDU_LEN_LT_6.hexlabeled", gen_G4_rows),
        ("G5_LEN_NEQ_7_PLUS_2Q.hexlabeled", gen_G5_rows),
        ("G6_QTY_OOB.hexlabeled", gen_G6_rows),
        ("G7_BC_NEQ_2Q.hexlabeled", gen_G7_rows),
        ("G8_ADDR_OOB.hexlabeled", gen_G8_rows),
    ]
    for fname, fn in specs:
        path = outdir / fname
        count = write_labeled(path, fn())
        print(f"[WRITE] {path} ({count} lines)")
    print("[OK] Horn rules =>", horn_path)
    print("[OK] HEX files  =>", outdir)

if __name__ == "__main__":
    main()